package com.idanchuang.component.sentinel;

import com.alibaba.csp.sentinel.adapter.spring.webmvc.callback.BlockExceptionHandler;
import com.alibaba.csp.sentinel.adapter.spring.webmvc.callback.RequestOriginParser;
import com.alibaba.csp.sentinel.adapter.spring.webmvc.config.SentinelWebMvcConfig;
import com.idanchuang.component.base.JsonResult;
import com.idanchuang.component.base.exception.common.ErrorCode;
import com.idanchuang.component.base.exception.core.ExFactory;
import com.idanchuang.component.core.Dc;
import com.idanchuang.component.core.helper.EnvContext;
import com.idanchuang.component.sentinel.config.SentinelConfigBean;
import com.idanchuang.component.sentinel.interceptor.ComponentSentinelWebInterceptor;
import com.idanchuang.component.sentinel.runner.SentinelConfigRunner;
import com.idanchuang.component.sentinel.runner.SwitchListenerRunner;
import com.idanchuang.spi.config.IConfig;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;

import java.io.PrintWriter;

/**
 * @author yjy
 * @date 2020/3/31 16:33
 **/
@Configuration
@Import(ComponentSentinelWebMvcConfig.class)
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
@ConditionalOnProperty(name = "spring.cloud.sentinel.enabled", matchIfMissing = true)
public class ComponentSentinelAutoConfiguration {

    @Bean
    @ConfigurationProperties(prefix = "idanchuang.component.sentinel")
    public SentinelConfigBean sentinelConfigBean() {
        return new SentinelConfigBean();
    }

    @Bean
    @ConditionalOnMissingBean(BlockExceptionHandler.class)
    public BlockExceptionHandler blockExceptionHandler(SentinelConfigBean sentinelConfigBean) {
        return (request, response, e) -> {
            response.setStatus(sentinelConfigBean.getResponseCode());
            response.setCharacterEncoding("UTF-8");
            response.setHeader("Content-type", "application/json;charset=UTF-8");
            PrintWriter out = response.getWriter();
            Object sentinelMsg = request.getAttribute(SentinelConfigBean.SENTINEL_MSG_KEY);
            if (sentinelMsg != null) {
                out.print(JsonResult.failure((String)sentinelMsg).toJsonString());
            } else {
                sentinelMsg = e.getMessage();
                if (sentinelMsg == null) {
                    sentinelMsg = e.getClass().getSimpleName();
                }
                out.print(JsonResult.failure(ExFactory.of(ErrorCode.TOO_MANY_REQUEST, sentinelMsg)).toJsonString());
            }
            out.flush();
            out.close();
        };
    }

    @Bean
    @ConditionalOnMissingBean(RequestOriginParser.class)
    public RequestOriginParser requestOriginParser(SentinelConfigBean sentinelConfigBean) {
        return request -> {
            // 优先获取 originReq
            if (sentinelConfigBean.isPreferOriginReq()) {
                String originReq = EnvContext.getAttribute(Dc.ORIGIN_REQUEST);
                if (originReq != null) {
                    return originReq;
                }
            }
            String originParser = sentinelConfigBean.getOriginParser();
            String originParseFrom = sentinelConfigBean.getOriginParseFrom();
            if (originParser == null || "".equals(originParser)) {
                originParser = "Origin";
            }
            String finalOriginParser = originParser;
            // 从参数中取
            if ("query".equalsIgnoreCase(originParseFrom)
                    || "parameter".equalsIgnoreCase(originParseFrom)) {
                return request.getParameter(finalOriginParser);
            }
            // 从header中取
            else if ("header".equalsIgnoreCase(originParseFrom)) {
                return request.getHeader(finalOriginParser);
            }
            // 默认从header中取
            return request.getHeader(finalOriginParser);
        };
    }

    @Bean
    @ConditionalOnMissingBean(SentinelWebMvcConfig.class)
    public SentinelWebMvcConfig sentinelWebMvcConfig(SentinelConfigBean sentinelConfigBean,
                                                     BlockExceptionHandler blockExceptionHandler,
                                                     RequestOriginParser requestOriginParser) {
        SentinelWebMvcConfig config = new SentinelWebMvcConfig();
        config.setBlockExceptionHandler(blockExceptionHandler);
        config.setHttpMethodSpecify(sentinelConfigBean.isHttpMethodSpecify());
        config.setOriginParser(requestOriginParser);
        return config;
    }

    @Bean
    @ConditionalOnMissingBean(ComponentSentinelWebInterceptor.class)
    public ComponentSentinelWebInterceptor sentinelWebInterceptor(SentinelWebMvcConfig sentinelWebMvcConfig) {
        return new ComponentSentinelWebInterceptor(sentinelWebMvcConfig);
    }

    @Bean("acgSentinelConfigRunner")
    @ConditionalOnMissingBean(name = "acgSentinelConfigRunner")
    public SentinelConfigRunner acgSentinelConfigRunner() {
        return new SentinelConfigRunner();
    }

    @Bean("acgSwitchListenerRunner")
    @ConditionalOnBean(IConfig.class)
    public SwitchListenerRunner acgSwitchListenerRunner() {
        return new SwitchListenerRunner();
    }

}
